#pragma once
#include "../common/net.h"
#include "../common/message.h"
#include <set>

namespace rpc
{
    namespace server
    {
        class ProviderManager
        {
        public:
            using ptr = std::shared_ptr<ProviderManager>;
            struct provider
            {
                using ptr = std::shared_ptr<provider>;
                std::mutex _mutex;
                BaseConnection::ptr conn;
                Address host;
                std::vector<std::string> _method;
                provider(const BaseConnection::ptr &c, Address h)
                    : conn(c), host(h)
                {
                }
                void addMethod(const std::string &method)
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    _method.emplace_back(method);
                }
            };
            void addProvider(const BaseConnection::ptr &conn, const Address &host, const std::string &method)
            {
                provider::ptr pro;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it = _conn.find(conn);
                    if (it == _conn.end())
                    {
                        pro = it->second;
                    }
                    else
                    {
                        pro = std::make_shared<provider>(conn, host);
                        _conn.insert(std::make_pair(conn, pro));
                    }
                    auto provides = _providers[method];
                    provides.insert(pro);
                }
                pro->addMethod(method);
            }
            // 当一个服务提供者断开连接的时候，获取他的信息--用于进行服务下线通知
            provider::ptr getProvider(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conn.find(conn);
                if (it != _conn.end())
                {
                    return it->second;
                }
                return provider::ptr();
            }
            void delProvider(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conn.find(conn);
                if (it == _conn.end())
                {
                    return;
                }
                for (auto &method : it->second->_method)
                {
                    auto &provider = _providers[method];
                    provider.erase(it->second);
                }
                _conn.erase(it);
            }
            std::vector<Address> methodHots(const std::string &method)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _providers.find(method);
                if (it == _providers.end())
                {
                    return std::vector<Address>();
                }
                std::vector<Address> result;
                for (auto &provider : it->second)
                {
                    result.push_back(provider->host);
                }
                return result;
            }

        private:
            std::mutex _mutex;
            std::unordered_map<std::string, std::set<provider::ptr>> _providers;
            std::unordered_map<BaseConnection::ptr, provider::ptr> _conn;
        };

        class DiscoverManager
        {
        public:
            using ptr = std::shared_ptr<DiscoverManager>;
            struct Discover
            {
                using ptr = std::shared_ptr<Discover>;
                std::mutex _mutex;
                BaseConnection::ptr conn;
                std::vector<std::string> _method;
                Discover(const BaseConnection::ptr &c)
                    : conn(c)
                {
                }
                void appendMethod(const std::string &method)
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    _method.push_back(method);
                }
            };
            Discover::ptr addDiscover(const BaseConnection::ptr &conn, std::string method)
            {
                Discover::ptr discover;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it = _conn.find(conn);
                    if (it != _conn.end())
                    {
                        discover = it->second;
                    }
                    else
                    {
                        discover = std::make_shared<Discover>(conn);
                        _conn.insert(std::make_pair(conn, discover));
                    }
                }
                discover->appendMethod(method);
                return discover;
            }
            void delDiscover(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conn.find(conn);
                if (it == _conn.end())
                {
                    return;
                }
                for (auto &method : it->second->_method)
                {
                    auto discovers = _discovers[method];
                    discovers.erase(it->second);
                }
                _conn.erase(it);
            }
            void onlineNotify(const std::string &method, const Address &host)
            {
                return notify(method, host, ServiceOptype::SERVICE_ONLINE);
            }
            // 当有一个服务提供者断开连接，则进行下线通知
            void offlineNotify(const std::string &method, const Address &host)
            {
                return notify(method, host, ServiceOptype::SERVICE_OFFLINE);
            }

        private:
            void notify(const std::string &method, const Address &host, ServiceOptype optype)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _discovers.find(method);
                if (it == _discovers.end())
                {
                    return;
                }
                auto msg_req = MessageFactory::create<ServiceRequest>();
                msg_req->setrId(UUID::uuid());
                msg_req->setMtype(MType::REQ_SERVICE);
                msg_req->Setmethod(method);
                msg_req->Setoptype(optype);
                msg_req->Sethost(host);
                for (auto &discover : it->second)
                {
                    discover->conn->send(msg_req);
                }
            }

        private:
            std::mutex _mutex;
            std::unordered_map<std::string, std::set<Discover::ptr>> _discovers;
            std::unordered_map<BaseConnection::ptr, Discover::ptr> _conn;
        };

        class PDManager
        {
        public:
            using ptr = std::shared_ptr<PDManager>;
            PDManager() : _provider(std::make_shared<ProviderManager>()), _discover(std::make_shared<DiscoverManager>())
            {
            }
            void onServiceRequest(const BaseConnection::ptr &conn, const ServiceRequest::ptr &req)
            {
                ServiceOptype optype = req->optype();
                if (optype == ServiceOptype::SERVICE_REGISTRY)
                {
                    ILOG("%s:%d 注册服务 %s", req->host().first.c_str(), req->host().second, req->method().c_str());
                    _provider->addProvider(conn, req->host(), req->method());
                    _discover->onlineNotify(req->method(), req->host());
                }
                else if (optype == ServiceOptype::SERVICE_DISCOVERY)
                {
                    ELOG("客户端要进行 %s 服务发现！", req->method().c_str());
                    _discover->addDiscover(conn, req->method());
                    return discoveryResponse(conn, req);
                }
                else
                {
                    ELOG("收到服务操作请求，但是操作类型错误！");
                    return errorResponse(conn, req);
                }
            }
            void onConnShutdown(const BaseConnection::ptr conn)
            {
                auto provider = _provider->getProvider(conn);
                if (provider.get() != nullptr)
                {
                    ILOG("%s:%d 服务下线", provider->host.first.c_str(), provider->host.second);
                    for (auto &method : provider->_method)
                    {
                        _discover->offlineNotify(method, provider->host);
                    }
                    _provider->delProvider(conn);
                }
                _discover->delDiscover(conn);
            }

        private:
            void errorResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                auto msg_rsp = MessageFactory::create<ServiceReponse>();
                msg_rsp->setrId(msg->rid());
                msg_rsp->setMtype(MType::RSP_SERVICE);
                msg_rsp->SetRCode(RCode::RCODE_INVALID_OPTYPE);
                msg_rsp->Setoptype(ServiceOptype::SERVICE_UNKNOW);
                conn->send(msg_rsp);
            }
            void registryResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                auto msg_rsp = MessageFactory::create<ServiceReponse>();
                msg_rsp->setrId(msg->rid());
                msg_rsp->setMtype(MType::RSP_SERVICE);
                msg_rsp->SetRCode(RCode::RCODE_OK);
                msg_rsp->Setoptype(ServiceOptype::SERVICE_REGISTRY);
                conn->send(msg_rsp);
            }
            void discoveryResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                auto msg_rsp = MessageFactory::create<ServiceReponse>();
                msg_rsp->setrId(msg->rid());
                msg_rsp->setMtype(MType::RSP_SERVICE);
                msg_rsp->Setoptype(ServiceOptype::SERVICE_DISCOVERY);
                std::vector<Address> hosts = _provider->methodHots(msg->method());
                if (hosts.empty())
                {
                    msg_rsp->SetRCode(RCode::RCODE_NOT_FOUND_SERVICE);
                    return conn->send(msg_rsp);
                }
                msg_rsp->SetRCode(RCode::RCODE_OK);
                msg_rsp->Setmethod(msg->method());
                msg_rsp->Sethost(hosts);
                return conn->send(msg_rsp);
            }

        private:
            ProviderManager::ptr _provider;
            DiscoverManager::ptr _discover;
        };
    }
}
